from elasticsearch import Elasticsearch
from Logger import Logger
from ConfigManager import ConfigManager
from typing import Dict, Any, List, Union


logger = Logger.get_logger(__name__)

class ElasticsearchClient:
    """Elasticsearch client wrapper for RAG application"""
    
    def __init__(self, url: str = "http://localhost:9200"):
        self.client = Elasticsearch(url, verify_certs=False,
            ssl_show_warn=False)
        self.index_name = ConfigManager().get('es_index_name')
        
    def create_index(self) -> bool:
        """Create the index with predefined mapping"""
        if self.client.indices.exists(index=self.index_name):
            return True
            
        mapping = {
            "mappings": {
                "properties": {
                    "doc_id": {"type": "keyword"},
                    "page_number": {"type": "text"},
                    "text_content": {
                        "type": "text",
                        "analyzer": "ik_max_word",
                        "search_analyzer": "ik_smart"
                    },
                    "doc_embedding": {
                        "type": "dense_vector",
                        "dims": 1024,
                        "index": True,
                        "similarity": "cosine"
                    },
                    "file_path": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword"}}
                    },
                    "file_size": {"type": "long"},
                    "file_format": {"type": "keyword"},
                    "last_modified": {
                        "type": "date",
                        "format": "strict_date_optional_time||epoch_millis"
                    }
                }
            }
        }
        
        try:
            self.client.indices.create(index=self.index_name, body=mapping,
                ignore=400)
            return True
        except Exception as e:
            logger.error(f"创建索引失败: {str(e)}")
            return False
            
    def index_document(self, document: Dict[str, Any]) -> bool:
        """Index a single document"""
        try:
            self.client.index(index=self.index_name, body=document)
            return True
        except Exception as e:
            logger.error(f"索引文档失败: {str(e)}")
            return False

    def vector_search(self, embedding: List[float], top_k: int = 5) -> List[Dict]:
        """Search documents by vector similarity"""
        try:
            query = {
                "size": top_k,
                "query": {
                    "script_score": {
                        "query": {"match_all": {}},
                        "script": {
                            "source": "cosineSimilarity(params.query_vector, 'doc_embedding') + 1.0",
                            "params": {"query_vector": embedding}
                        }
                    }
                }
            }
            response = self.client.search(index=self.index_name, body=query)
            return [hit["_source"] for hit in response["hits"]["hits"]]
        except Exception as e:
            logger.error(f"向量搜索失败: {str(e)}")
            return []


